package test.demo;


import com.alibaba.fastjson.JSON;
import tech.waterism.modelbase.PredictBase;
import tech.waterism.modelbase.PredictResult;

import java.util.*;

/**
 * 陕北模型 超渗产流 三水源
 * @author 初京刚
 * @date 2022-3-4
 * @version 1.0
 */

public class SHANNXI_R_3 extends PredictBase {

    /**
     * 流域平均张力水容量（mm）、流域蓄水容量
     */
    private double Wm;

    /**
     * 自由水水深
     */
    private double sp;

    /**
     * 产流面积系数
     */
    private double frp;

    /**
     * 流域蒸散发折算系数
     */
    private double Kc;

    /**
     * 流域内下渗能力抛物线形状系数
     */
    private double Bx;

    /**
     * 不透水的面积占全流域面积的比
     */
    private double IM;

    /**
     * 流域自由水蓄水容量
     */
    private double SM;

    /**
     * 自由水蓄水容量曲线指数
     */
    private double EX;

    /**
     * 地下水径流出流系数（cjg）表层自由水蓄水库对地下水的日出流系数--有Kg+Ki=0.7
     */
    private double KG;

    /**
     * 壤中水径流出流系数
     */
    private double KI;

    /**
     * 流域 月蒸发量（mm）
     */
    private double[] ES;

    /**
     * 流域初始下渗能力，相当于土壤干燥时的下渗能力，平均最大下渗能力，mm/min
     */
    private double f0;

    /**
     * 流域平均稳定下渗能力，mm/min
     */
    private double fc;

    /**
     * 土壤含水计算的允许误差
     */
    private double err_limit;

    /**
     * 下渗能力衰弱系数，随土质而变的系数，1/min
     */
    private double k;

    /**
     * 流域 时段蒸发量（mm）
     */
    private double[] floodE;

    /**
     * 降水时间
     */
    private Date[] floodTm;

    /**
     * 降水量
     */
    private double[] floodDrp;

    /**
     * 降水序列长度
     */
    private Integer floodRainRange;

    /**
     * 降水时段长
     */
    private double timeInterval;

    /**
     * 模拟总产流
     */
    private double[] RunoffSim;

    /**
     * 总产流
     */
    private double[] R;

    /**
     * 分水源-地表水
     */
    private double[] RS;

    /**
     * 分水源-壤中水
     */
    private double[] RI;

    /**
     * 分水源-地下水
     */
    private double[] RG;

//    /**
//     * 不透水面积产流
//     */
//    private double[] RIm;

    /**
     * 净雨
     */
    private double[] PE;

    /**
     * 雨强 mm/min
     */
    private double[] PEI;

    /**
     * 土壤含水量
     */
    private double[] Ws;

    /**
     * 产流面积自由水水深
     */
    private double[] S;

    /**
     * 产流面积系数
     */
    private double[] FR;

    /**
     * 流域面积
     */
    private double Area;

    /**
     * 流域张力水含量初值（m）
     */
    private double WSInitial;

    /**
     * 下渗曲线选取标志
     */
    private int type;

    /**
     * 菲利普下渗曲线参数
     */
    private double A;

    /**
     * 菲利普下渗曲线参数
     */
    private double B;

    public SHANNXI_R_3(Map<String, String> data) {
        super();
        WSInitial = Double.parseDouble(data.get("WSIG"));
        sp = Double.parseDouble(data.get("SP"));
        frp = Double.parseDouble(data.get("FRP"));
        Kc = Double.parseDouble(data.get("KC"));
        Bx = Double.parseDouble(data.get("BX"));
        IM = Double.parseDouble(data.get("IM"));
        type = Integer.parseInt(data.get("TYPE"));
        f0 = Double.parseDouble(data.get("F0"));
        fc = Double.parseDouble(data.get("FC"));
        err_limit = Double.parseDouble(data.get("ERLMT"));
        k = Double.parseDouble(data.get("K"));
        Wm = f0 / k;
        A = Double.parseDouble(data.get("A"));
        B = Double.parseDouble(data.get("B"));
        SM = Double.parseDouble(data.get("SM"));
        EX = Double.parseDouble(data.get("EX"));
        KG = Double.parseDouble(data.get("KG"));
        KI = Double.parseDouble(data.get("KI"));
        ES = JSON.parseObject(data.get("ES"), double[].class);
        floodTm = JSON.parseObject(data.get("dt"), Date[].class);// input[i].data.dt
        floodDrp = JSON.parseObject(data.get("rain"), double[].class);// input[i].data.rain
        floodE = JSON.parseObject(data.get("evaporation"), double[].class);  // input[i].data.evaporation
        floodRainRange = floodDrp.length;
        timeInterval = Double.parseDouble(data.get("clen"));                  // index.clen
        RunoffSim = new double[floodRainRange];
        R = new double[floodRainRange];
//        RIm = new double[floodRainRange];
        PE = new double[floodRainRange];
        PEI = new double[floodRainRange];
        RS = new double[floodRainRange];
        RI = new double[floodRainRange];
        RG = new double[floodRainRange];
    }

    /**
     * 模型计算
     */
    public Map<String, double[]> SHANNXImodel() {

        // 超渗产流
        expinfMechanism();

        // 分水源
        freeTank();

        for (Integer j = 0; j < floodRainRange; j++) {
//            RunoffSim[j] = RS[j] + RIm[j] + RI[j] + RG[j];
            RunoffSim[j] = RS[j] + RI[j] + RG[j];
        }

        Map<String, double[]> simulationResult = new HashMap<>();
        simulationResult.put("rain", floodDrp);
        simulationResult.put("PE", PE);
        simulationResult.put("R", R);
        simulationResult.put("runoffSim", RunoffSim);
        simulationResult.put("RS", RS);
        simulationResult.put("RI", RI);
        simulationResult.put("RG", RG);
//        simulationResult.put("RIM", RIm);

       return simulationResult;
    }

    /**
     * 蒸发计算
     *
     * @param date   降水时间
     * @param timeIn 时间间隔 (h)
     * @return 蒸发
     */
    private double calcEvaporationPotential(Date date, double timeIn) {
        Calendar time = Calendar.getInstance();
        time.setTime(date);
        Integer Month = time.get(Calendar.MONTH) + 1; // 获得降雨开始月份
        double Ep = ES[Month - 1]/30; // 月蒸发量转换成日蒸发量
        Ep = Ep * timeIn / 24; // 日蒸发量转换成时段蒸发量
        return Ep;
    }

    /**
     * 超渗产流模型
     *
     * @return 总产流
     */
    private void expinfMechanism() {

        Ws = new double[floodRainRange + 1]; // 张力水蓄量，随时段变化
        Ws[0] = WSInitial;

        double EPt; // EPt为蒸散发，等于Kc * E（实测蒸发，即ES[]或时段）
        double Ft0 = 0.0; // 时段初流域下渗能力
        double Ft = 0.0; //时段下渗量
        double fmm;

        for(Integer i = 0; i < floodRainRange; i++) {

            //蒸发计算
            if (floodE == null || floodE.length != floodRainRange) {
                EPt = Kc * calcEvaporationPotential(floodTm[i], timeInterval);
            } else {
                EPt = Kc * floodE[i];
            }

            // PE[i]为降雨减去蒸发，可正可负
            PE[i] = floodDrp[i] - EPt;
            PEI[i] = PE[i] / (60 * timeInterval);

            if (type == 0) {
                //霍顿下渗曲线<霍顿与菲利蒲下渗公式对子洲径流站资料的拟合>
                Ft0 = HortonInfil(Ws[i]);
                // 计算流域该时段的最大点下渗能力<格林一安普特下渗曲线的改进和应用>、<格林－安普特降雨径流模型改进及初步应用>
                fmm = Ft0 * (1 + Bx) / (1 - IM);
                if (PEI[i] > fmm) {
                    // 全流域产流
                    /*Ws[i + 1] = fc * 60 * timeInterval * (i + 1) + 1 / k * (f0 - fc) * (1 - Math.exp(-k * 60 * timeInterval * (i + 1)));
                    if(Ws[i + 1] > Wm){
                        Ws[i + 1] = Wm;
                    }
                    Ft = Ws[i + 1] - Ws[i];*/
                    Ft = Ft0 * 60 * timeInterval;
                    Ws[i + 1] = Ws[i] + Ft;
                    if(Ft < 0){
                        Ft = 0;
                    }
                    if(Ws[i + 1] > Wm){
                        Ws[i + 1] = Wm;
                    }
                    if(Ws[i + 1] < 0){
                        Ws[i + 1] = 0;
                    }
                    R[i] = PE[i] - Ft;
                    if(R[i] < 0){
                        R[i] = 0;
                    }
                } else {
                    R[i] = (PEI[i] - Ft0 + Ft0 * Math.pow(1 - PEI[i] / fmm, Bx + 1)) * 60 * timeInterval;
                    if(R[i] < 0){
                        R[i] = 0;
                    }
                    Ws[i + 1] = Ws[i] + PE[i] - R[i];
                    if(Ws[i + 1] > Wm){
                        Ws[i + 1] = Wm;
                    }
                    if(Ws[i + 1] < 0){
                        Ws[i + 1] = 0;
                    }
                }
            } else if (this.type == 1) {
                //菲利普斯下渗曲线<霍顿与菲利蒲下渗公式对子洲径流站资料的拟合>
                if(Ws[i] > 0) {
                    Ft0 = PhillipsInfli(Ws[i]);
                } else {
                    Ft0 = PhillipsInfli(0.1);
                }
                // 计算流域该时段的最大点下渗能力<格林一安普特下渗曲线的改进和应用>、<格林－安普特降雨径流模型改进及初步应用>
                fmm = Ft0 * (1 + Bx) / (1 - IM);
                if (PEI[i] > fmm) {
                    // 全流域产流
                    R[i] = (PEI[i] - Ft0) * 60 * timeInterval;
                    if(R[i] < 0){
                        R[i] = 0;
                    }
                    Ws[i + 1] = Ws[i] + Ft0 * 60 * timeInterval;
                    if(Ws[i + 1] > Wm){
                        Ws[i + 1] = Wm;
                    }
                } else {
                    R[i] = (PEI[i] - Ft0 + Ft0 * Math.pow(1 - PEI[i] / fmm, Bx + 1)) * 60 * timeInterval;
                    if(R[i] < 0){
                        R[i] = 0;
                    }
                    Ws[i + 1] = Ws[i] + PE[i] - R[i];
                    if(Ws[i + 1] > Wm){
                        Ws[i + 1] = Wm;
                    }
                    if(Ws[i + 1] < 0){
                        Ws[i + 1] = 0;
                    }
                }
            }
        }

    }

    /**
     * 菲利普斯下渗曲线计算
     *
     * @param w0          时段初土壤含水
     * @return 下渗能力
     */
    private double PhillipsInfli(double w0) {
        return B * B * (1 + Math.sqrt(1 + 4 * A * w0 / (B * B))) / (4 * w0) + A;
    }

    /**
     * 霍顿下渗曲线计算
     *
     * @param w0          时段初土壤含水
     * @return 下渗能力
     */
    private double HortonInfil(double w0) {

        /*
            理论基础：<流域水文模型-新安江模型和陕北模型>pdf66页
            数学基础：迭代求解
        */
        // 由经验初设模型参数
        double T = w0 / f0;

        double w0ByCount = fc * T + 1 / k * (f0 - fc) * (1 - Math.exp(-k * T));
        double fByCount = 0.0;
        while (Math.abs(w0 - w0ByCount) > err_limit) {
            fByCount = f0 - k * (w0ByCount - fc * T);
//            fByCount = fc + (f0 - fc) * Math.exp(-k * T);
            T = T + (w0 - w0ByCount) / fByCount;
            w0ByCount = fc * T + 1 / k * (f0 - fc) * (1 - Math.exp(-k * T));
        }
        fByCount = f0 - k * (w0 - fc * T);
        return fByCount;
    }

    /**
     * 分水源模型
     *
     * @return 地表水、壤中水、地下水
     */
    private void freeTank() {
        S = new double[floodRainRange + 1];// 存每一时段初产流面积上的自由水蓄量深
        FR = new double[floodRainRange + 1];// 存每一时段的产流面积系数（认为时段内的产流面积不变）
//        S[0] = 0.0;// 假定初始自由水蓄量为0
//        FR[0] = 0.001; // 假定初始产流面积为0.001
        S[0] = sp; // 设定产流面积初始自由水蓄量
        FR[0] = frp; // 设定初始产流面积系数

        // 参数KI和KG通常以日尺度定义，需要转换为时段尺度
        double D = 24 / timeInterval;
        double KID = (1 - Math.pow((1 - (KI + KG)), 1.0 / D)) / (1 + KG / KI);
        double KGD = KID * KG / KI;

        double EPt; // 蒸散发，等于Kc * E（实测蒸发，即ES[]或时段）
        double PEt; // PEt为净雨，等于降雨减去蒸发

        for (Integer i = 0; i < floodRainRange; i++) {

            // 蒸发计算
            if(floodE==null || floodE.length!=floodRainRange){
                EPt = Kc * calcEvaporationPotential(floodTm[i], timeInterval);
            } else{
                EPt = Kc * floodE[i];
            }

            // PEt为降雨减去蒸发，可正可负
            PEt = floodDrp[i] - EPt;

            // 范围控制
            S[i] = S[i] < 0 ? 0 : S[i];

            // 有产流量才需要分段计算
            if (PEt > 0) {

                // 通过径流系数计算产流面积,产流面积等于流域总径流量/净雨
                FR[i] = R[i] / PEt;

                // 此处一定要控制产流面积的范围，因为计算出的产流面积有可能不在合理范围内
                if (FR[i] > 1.0)
                    // 注意产流面积不能比1大
                    FR[i] = 1.0;
                if (FR[i] < 0.0)
                    // 注意产流面积不能比0小
                    FR[i] = 0.0;
                if (i != 0) {
                    // 由于产流面积的变化，需将上个时段末的自由水蓄量换算成本时段初对应产流面积上的深度值）
                    S[i] = S[i] * FR[i - 1] / FR[i];
                }

                // 以5mm分段，划分成多段进行计算
                double q = R[i] / FR[i]; // 产流面积上净雨
                double N = Math.floor(q / 5 + 1);
                double pe = q / N;// 每段的净雨深

                // 时段长改变，参数也要相应变化
                double KIDN = (1 - Math.pow((1 - (KID + KGD)), 1.0 / N)) / (1 + KGD / KID);
                double KGDN = KIDN * KGD / KID;

                double SMM = SM * (1 + EX);// 流域单点最大自由水蓄水容量

                // 产流面积上的单点最大自由水蓄水容量
                double SMMF;

                if (EX <= 0.0001) {
                    SMMF = SMM;
                } else {
                    SMMF = SMM * (1.0 - Math.pow(1.0 - FR[i], 1.0 / EX));
                }

                double SMF = SMMF / (1.0 + EX); //产流面积上平均自由水蓄水容量

                double s = S[i]; // 该时段初产流面积上的自由水蓄量
                double rs = 0.0; // 该时段流域上的地表产流
                double ri = 0.0; // 该时段流域上的土壤产流
                double rg = 0.0; // 该时段流域上的地下产流
                for (Integer j = 0; j < N; j++) {
                    // 若自由水蓄水库本身已蓄满
                    if (s >= SMF) {
                        rs = pe * FR[i];
                        ri = SMF * KIDN * FR[i];
                        rg = SMF * KGDN * FR[i];
                        s = SMF - (ri + rg) / FR[i];
                    } else {
                        double AU = SMMF * (1 - Math.pow(1 - s / SMF, 1 / (1 + EX)));
                        if (pe + AU < SMMF) {
                            rs = FR[i] * (pe - SMF + s + SMF * Math.pow(1 - (pe + AU) / SMMF, 1 + EX));
                            double SMt = pe + s - rs / FR[i]; //时段产流面积上扣除地表产流后的自由水蓄量
                            ri = SMt * KIDN * FR[i];
                            rg = SMt * KGDN * FR[i];
                            s += pe - (rs + ri + rg) / FR[i];
                        } else {
                            rs = (pe + s - SMF) * FR[i];
                            ri = SMF * KIDN * FR[i];
                            rg = SMF * KGDN * FR[i];
                            s = SMF - (ri + rg) / FR[i];
                        }
                    }
                    RS[i] += rs;
                    RI[i] += ri;
                    RG[i] += rg;
                }

                /*需保留最后一个时段末的状态值，去掉时段初值限制条件*/
//                if(i + 1 < floodRainRange){
                S[i + 1] = s;
//                }
            } else {
                // 无产流量，产流面积与上个时段一样
                if (i != 0) {
                    FR[i] = FR[i - 1];
                    // 此处一定要控制产流面积的范围，因为计算出的产流面积有可能不在合理范围内
                    if (FR[i] > 1.0)
                        // 注意产流面积不能比1大
                        FR[i] = 1.0;
                    if (FR[i] < 0.0)
                        // 注意产流面积不能比0小
                        FR[i] = 0.0;
                }
                RS[i] = 0.0;
                RI[i] = KID * S[i] * FR[i];
                RG[i] = KGD * S[i] * FR[i];

                /*需保留最后一个时段末的状态值，去掉时段初值限制条件*/
//                if(i + 1 < floodRainRange){
                S[i + 1] = S[i] * (1 - KID - KGD);

                if (S[i + 1] < 0.0)
                    S[i + 1] = 0.0;
                if (S[i + 1] > SM)
                    S[i + 1] = SM;
//                }
            }
        }
    }


    private List<String> toStringList(double[] array0){
        String[] array = new String[array0.length];
        for(int i = 0; i < array.length; i++){
            array[i] = String.valueOf(array0[i]);
        }
        return Arrays.asList(array);
    }

    private List<String> toStringList(double[][] array0){
        String[] array = new String[array0.length];
        for(int i = 0; i < array.length; i++){
            array[i] = Arrays.toString(array0[i]);
        }
        return Arrays.asList(array);
    }

    @Override
    public void init() {

    }

    @Override
    public PredictResult predict() {
        PredictResult result = new PredictResult();
        Map<String, double[]> map = SHANNXImodel();

        double[][] RSim = new double[4][];
        //总产流
        RSim[0] = map.get("R");
        //流域地表产流
        RSim[1] = map.get("RS");
        //流域土壤产流
        RSim[2] = map.get("RI");
        //流域地下产流
        RSim[3] = map.get("RG");

        //流域面积总产流
        result.setRunoffSim(map.get("runoffSim"));
        result.setRSim(RSim);

        // 更新状态
        List<String> WS = toStringList(Ws);
//        List<String> WUP = toStringList(Wu);
//        List<String> WLP = toStringList(Wl);
//        List<String> WDP = toStringList(Wd);
        List<String> SP = toStringList(S);
        List<String> FRP = toStringList(FR);
//        List<String> QXSIG = toStringList(QX0);
//        List<String> QSIG = toStringList(QS0);
//        List<String> QSP = toStringList(QS);
//        List<String> QIP = toStringList(QIS0);
//        List<String> QGP = toStringList(QGS0);
//        List<String> QSP = toStringList(QS);
//        List<String> QIP = toStringList(QI);
//        List<String> QGP = toStringList(QG);
        Map<String, List<String>> status = new HashMap<>();
        //土壤含水量
        status.put("WS", WS);
//        //上层张力水含量
//        status.put("WUP", WUP);
//        //下层张力水含量
//        status.put("WLP", WLP);
//        //深层张力水含量
//        status.put("WDP", WDP);
        //产流面积自由水深
        status.put("SP", SP);
        //产流面积系数
        status.put("FRP", FRP);
//        //河段下断面各时段初始流量
//        status.put("QXSIG", QXSIG);
//        //各时段初始总基流
//        status.put("QSIG", QSIG);
//        //地表坡面汇流
//        status.put("QSP", QSP);
//        //坡面汇流-壤中流
//        status.put("QIP", QIP);
//        //地下坡面汇流
//        status.put("QGP", QGP);
        result.setNewStatus(status);

        return result;

        //return null;
    }
}
